import socket
from scapy.all import ARP, Ether, srp
import subprocess
import platform
import re
import json
from typing import List, Dict, Optional, Tuple
import concurrent.futures
import threading
from functools import lru_cache
import time
import paramiko

COMMON_PORTS = {
    20: "ftp-data",
    21: "ftp",
    22: "ssh",
    23: "telnet",
    25: "smtp",
    53: "dns",
    67: "dhcp-server",
    68: "dhcp-client",
    69: "tftp",
    80: "http",
    110: "pop3",
    111: "rpcbind",
    123: "ntp",
    135: "msrpc",
    137: "netbios-ns",
    138: "netbios-dgm",
    139: "netbios-ssn",
    143: "imap",
    161: "snmp",
    162: "snmptrap",
    179: "bgp",
    389: "ldap",
    443: "https",
    445: "microsoft-ds",
    465: "smtps",
    514: "syslog",
    515: "printer",
    587: "smtp (submission)",
    631: "ipp",
    636: "ldaps",
    873: "rsync",
    993: "imaps",
    995: "pop3s",
    1080: "socks",
    1433: "mssql",
    1521: "oracle",
    1723: "pptp",
    1900: "upnp",
    2049: "nfs",
    2082: "cpanel",
    2083: "cpanel-ssl",
    2181: "zookeeper",
    2375: "docker-api",
    3306: "mysql",
    3389: "rdp",
    3690: "svn",
    4000: "remote-anything",
    4040: "webhook-dev",
    4369: "epmd",
    5000: "http-alt/flask",
    5001: "commplex-link",
    5432: "postgresql",
    5672: "amqp/rabbitmq",
    5900: "vnc",
    5984: "couchdb",
    6000: "x11",
    6379: "redis",
    6667: "irc",
    7001: "weblogic",
    7077: "spark",
    7474: "neo4j",
    8000: "http-dev",
    8008: "http-alt",
    8080: "http-proxy",
    8081: "http-alt-2",
    8443: "https-alt",
    8500: "consul",
    8888: "jupyter",
    9000: "php-fpm",
    9090: "prometheus",
    9200: "elasticsearch",
    9300: "elasticsearch-node",
    9999: "test-port",
    11211: "memcached",
    27017: "mongodb"
}

def get_service_name(port):
    return COMMON_PORTS.get(port, "unknown")

def scan_ports_verbose(host: str, ports: list[int], timeout=0.5):
    result = []
    for port in ports:
        try:
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
                s.settimeout(timeout)
                code = s.connect_ex((host, port))
                is_open = (code == 0)
        except Exception as e:
            is_open = False
        result.append({
            "port": port,
            "status": "open" if is_open else "closed",
            "service": get_service_name(port)
        })
    return result

def get_local_ip() -> str:
    """获取本机IP地址"""
    try:
        s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        s.connect(("8.8.8.8", 80))
        local_ip = s.getsockname()[0]
        s.close()
        return local_ip
    except:
        return "127.0.0.1"

@lru_cache(maxsize=100)
def get_hostname(ip: str) -> Optional[str]:
    """获取IP地址对应的主机名（带缓存）"""
    try:
        hostname = socket.gethostbyaddr(ip)[0]
        return hostname
    except (socket.herror, socket.gaierror):
        return None

def get_device_vendor(mac: str) -> Optional[str]:
    """根据MAC地址获取设备厂商信息"""
    try:
        # 移除MAC地址中的分隔符
        mac = mac.replace(':', '').replace('-', '').upper()
        # 获取OUI（前6位）
        oui = mac[:6]
        
        # 从本地缓存或在线数据库获取厂商信息
        # 这里可以添加本地MAC地址数据库查询
        # 或者调用在线API（需要网络连接）
        return None
    except:
        return None

def get_device_os(ip: str) -> Optional[str]:
    """尝试识别设备操作系统"""
    try:
        # 使用nmap进行操作系统检测（需要安装nmap）sudo apt-get install nmap
        result = subprocess.check_output(
            ['nmap', '-O', '--osscan-limit', ip],
            stderr=subprocess.STDOUT
        ).decode()
        
        # 解析nmap输出
        if "OS details:" in result:
            os_match = re.search(r"OS details: (.*?)(?:\n|$)", result)
            if os_match:
                return os_match.group(1).strip()
    except:
        pass
    return None

def get_open_ports(ip: str, ports: List[int] = None) -> List[Dict]:
    """获取设备开放的端口"""
    if ports is None:
        ports = [21, 22, 23, 25, 53, 80, 443, 445, 3306, 3389, 8080]
    
    open_ports = []
    for port in ports:
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(1)
            result = sock.connect_ex((ip, port))
            if result == 0:
                try:
                    service = socket.getservbyport(port)
                except:
                    service = "unknown"
                open_ports.append({
                    "port": port,
                    "service": service
                })
            sock.close()
        except:
            continue
    return open_ports

def get_device_info(ip: str, mac: str, local_ip: str) -> Dict:
    """获取单个设备的信息"""
    if ip == local_ip:
        return {
            "ip": ip,
            "hostname": socket.gethostname(),
            "type": "local",
            "status": "online",
            "os": platform.system() + " " + platform.release(),
            "open_ports": scan_ports_verbose(ip, list(COMMON_PORTS.keys())[:10])  # 只扫描前10个常用端口
        }
    
    # 使用线程池并发获取设备信息
    with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
        hostname_future = executor.submit(get_hostname, ip)
        ports_future = executor.submit(scan_ports_verbose, ip, list(COMMON_PORTS.keys())[:10])  # 只扫描前10个常用端口
        
        # 获取结果
        hostname = hostname_future.result(timeout=2)  # 2秒超时
        open_ports = ports_future.result(timeout=5)   # 5秒超时
    
    return {
        "ip": ip,
        "mac": mac,
        "hostname": hostname,
        "type": "network",
        "status": "online",
        "vendor": get_device_vendor(mac),
        "open_ports": open_ports
    }

def get_network_devices() -> Tuple[List[Dict], str]:
    """获取网络设备信息（优化版本）"""
    devices = []
    system = platform.system()
    local_ip = get_local_ip()
    
    # 添加本机信息
    devices.append(get_device_info(local_ip, "", local_ip))
    
    if system == "Darwin":  # macOS
        try:
            # 使用arp -a命令获取局域网设备
            output = subprocess.check_output(["arp", "-a"]).decode()
            device_futures = []
            
            # 使用线程池并发处理设备信息
            with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
                for line in output.splitlines():
                    if "incomplete" not in line:
                        match = re.search(r"\((\d+\.\d+\.\d+\.\d+)\) at ([0-9a-f:]+)", line)
                        if match:
                            ip = match.group(1)
                            mac = match.group(2)
                            if ip != local_ip:  # 排除本机
                                future = executor.submit(get_device_info, ip, mac, local_ip)
                                device_futures.append(future)
                
                # 获取所有设备信息（带超时）
                for future in concurrent.futures.as_completed(device_futures, timeout=30):
                    try:
                        device_info = future.result(timeout=5)
                        devices.append(device_info)
                    except Exception as e:
                        devices.append({
                            "error": f"获取设备信息失败: {str(e)}",
                            "type": "error"
                        })
                        
        except Exception as e:
            devices.append({
                "error": f"获取网络设备信息失败: {str(e)}",
                "type": "error"
            })
            
    elif system == "Linux":
        try:
            # 使用ip neigh命令获取局域网设备
            output = subprocess.check_output(["ip", "neigh", "show"]).decode()
            device_futures = []
            
            # 使用线程池并发处理设备信息
            with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
                for line in output.splitlines():
                    if "FAILED" not in line:
                        parts = line.split()
                        if len(parts) >= 4:
                            ip = parts[0]
                            mac = parts[4]
                            if ip != local_ip:  # 排除本机
                                future = executor.submit(get_device_info, ip, mac, local_ip)
                                device_futures.append(future)
                
                # 获取所有设备信息（带超时）
                for future in concurrent.futures.as_completed(device_futures, timeout=30):
                    try:
                        device_info = future.result(timeout=5)
                        devices.append(device_info)
                    except Exception as e:
                        devices.append({
                            "error": f"获取设备信息失败: {str(e)}",
                            "type": "error"
                        })
                        
        except Exception as e:
            devices.append({
                "error": f"获取网络设备信息失败: {str(e)}",
                "type": "error"
            })
    
    return devices, local_ip

class SSHConnection:
    def __init__(self):
        self.connections = {}
        self.working_dirs = {}  # 存储每个会话的工作目录
        self.lock = threading.Lock()

    def get_connection(self, session_id: str) -> Optional[paramiko.SSHClient]:
        with self.lock:
            return self.connections.get(session_id)

    def create_connection(self, session_id: str, host: str, port: int, username: str, password: str) -> Tuple[bool, Dict]:
        try:
            ssh = paramiko.SSHClient()
            ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            ssh.connect(
                hostname=host,
                port=port,
                username=username,
                password=password,
                timeout=10
            )
            with self.lock:
                self.connections[session_id] = ssh
                # 初始化工作目录
                stdin, stdout, stderr = ssh.exec_command('pwd')
                self.working_dirs[session_id] = stdout.read().decode('utf-8').strip()
            return True, {"status": "success", "message": "连接成功"}
        except Exception as e:
            return False, {"status": "error", "message": f"连接失败：{str(e)}"}

    def close_connection(self, session_id: str):
        with self.lock:
            if session_id in self.connections:
                self.connections[session_id].close()
                del self.connections[session_id]
            if session_id in self.working_dirs:
                del self.working_dirs[session_id]

    def execute_command(self, session_id: str, command: str) -> Tuple[bool, Dict]:
        ssh = self.get_connection(session_id)
        if not ssh:
            return False, {"status": "error", "message": "未找到活动连接"}

        try:
            # 如果是 cd 命令，更新工作目录
            if command.strip().startswith('cd '):
                # 获取当前工作目录
                current_dir = self.working_dirs.get(session_id, '')
                # 构建完整的 cd 命令
                if command.strip() == 'cd':
                    # 如果只是 cd，切换到用户主目录
                    full_command = 'cd ~ && pwd'
                else:
                    # 处理相对路径和绝对路径
                    target_dir = command.strip()[3:].strip()
                    if target_dir.startswith('/'):
                        # 绝对路径
                        full_command = f'cd {target_dir} && pwd'
                    else:
                        # 相对路径
                        full_command = f'cd {current_dir}/{target_dir} && pwd'
                
                # 执行命令
                stdin, stdout, stderr = ssh.exec_command(full_command, timeout=30)
                new_dir = stdout.read().decode('utf-8').strip()
                error = stderr.read().decode('utf-8')
                
                if not error:
                    with self.lock:
                        self.working_dirs[session_id] = new_dir
                    return True, {
                        "status": "success",
                        "output": "",
                        "error": "",
                        "exit_status": 0
                    }
                else:
                    return False, {
                        "status": "error",
                        "message": f"切换目录失败：{error}"
                    }
            else:
                # 对于非 cd 命令，在正确的工作目录下执行
                current_dir = self.working_dirs.get(session_id, '')
                # 使用完整的 shell 环境执行命令，确保命令被正确引用
                full_command = f'cd {current_dir} && source ~/.zshrc && {command}'
                
                # 使用 shell=True 来确保命令被正确解析
                stdin, stdout, stderr = ssh.exec_command(full_command, timeout=30)
                output = stdout.read().decode('utf-8')
                error = stderr.read().decode('utf-8')
                
                return True, {
                    "status": "success",
                    "output": output,
                    "error": error,
                    "exit_status": stdout.channel.recv_exit_status()
                }
        except Exception as e:
            return False, {"status": "error", "message": f"执行失败：{str(e)}"}

# 创建全局 SSH 连接管理器
ssh_manager = SSHConnection()

def execute_ssh_command(
    host: str,
    port: int,
    username: str,
    password: str,
    command: str = '',
    session_id: str = None,
    connection_timeout: int = 10,
    execution_timeout: int = 30
) -> Tuple[bool, Dict]:
    """
    执行SSH命令
    
    Args:
        host: 主机地址
        port: 端口号
        username: 用户名
        password: 密码
        command: 要执行的命令
        session_id: 会话ID
        connection_timeout: 连接超时时间
        execution_timeout: 执行超时时间
        
    Returns:
        Tuple[bool, Dict]: (是否成功, 结果字典)
    """
    if not session_id:
        # 如果没有会话ID，创建新连接
        success, result = ssh_manager.create_connection(
            session_id=str(time.time()),  # 使用时间戳作为临时会话ID
            host=host,
            port=port,
            username=username,
            password=password
        )
        if not success:
            return False, result
        
        # 执行命令
        success, result = ssh_manager.execute_command(session_id, command)
        # 关闭临时连接
        ssh_manager.close_connection(session_id)
        return success, result
    else:
        # 使用现有会话
        return ssh_manager.execute_command(session_id, command)

def get_system_info(ssh: paramiko.SSHClient, timeout: int) -> Dict:
    """
    获取远程系统信息
    
    Args:
        ssh: SSH客户端实例
        timeout: 执行超时时间
        
    Returns:
        Dict: 系统信息字典
    """
    system_info = {}
    
    # 获取系统基本信息
    stdin, stdout, stderr = ssh.exec_command('uname -a', timeout=timeout)
    system_info['system_info'] = stdout.read().decode('utf-8')
    
    # 获取磁盘信息
    stdin, stdout, stderr = ssh.exec_command('df -h', timeout=timeout)
    system_info['disk_info'] = stdout.read().decode('utf-8')
    
    # 获取内存信息
    stdin, stdout, stderr = ssh.exec_command('free -h', timeout=timeout)
    system_info['memory_info'] = stdout.read().decode('utf-8')
    
    # 获取CPU信息
    stdin, stdout, stderr = ssh.exec_command('top -bn1 | head -n 5', timeout=timeout)
    system_info['cpu_info'] = stdout.read().decode('utf-8')
    
    # 获取网络信息
    stdin, stdout, stderr = ssh.exec_command('netstat -tuln', timeout=timeout)
    system_info['network_info'] = stdout.read().decode('utf-8')
    
    # 获取进程信息
    stdin, stdout, stderr = ssh.exec_command('ps aux | head -n 10', timeout=timeout)
    system_info['process_info'] = stdout.read().decode('utf-8')
    
    # 获取系统负载
    stdin, stdout, stderr = ssh.exec_command('uptime', timeout=timeout)
    system_info['load_info'] = stdout.read().decode('utf-8')
    
    return system_info

